# -*- coding: utf-8 -*-
"""
#-------------------------------------------------------------------------------
# Nom:         optimGR
#
# Auteur:      Alain Gauthier
#
# Créé le:     23/08/2013
# Copyright:   (c) eaucéa 2014
# Licence:     GNU Public Licence (GPL)

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/

Module permettant l'optimisation des paramètres de GR selon des critères
spécifiés ('Nash' ou bien critère régional)

#-------------------------------------------------------------------------------
"""

# modules système
import datetime
# modules externes
import numpy as np
import numpy.linalg as LA
import scipy.optimize as opt
# modules internes
import analyseChroniques
import typeCritereNash
import gr2m

#-------------------------------------------------------------------------------

def lancerOptimisation(AcritereEfficacite,
                       AchroniquesPluie,
                       AchroniquesEvapo,
                       AchroniquesQmesure,
                       AdonneesStation,
                       AparamOptim,
                       fichierLog):
    """
    Optimisation des paramètres de GR4J ou GR2M pour la station de mesure.
    AcritereEfficacite : identifiant de la fonction Nash à optimiser
    AchroniquesPluie   : chronique de pluie sur le BV de la station
    AchroniquesEvapo   : chronique d'evapotranspiration sur le BV de la station
    AchroniquesQmesure : chronique des débits mesurés par la station de mesure
    AdonneesStation    : dict de paramètres nécessaires au calage (Xparam : valeurs initiales des paramètres X de GR4J/GR2M,
                         duree_eval : (debut,fin) intervalle de temps de l'évaluation du critère d'efficacité,
                         aire : aire du BV au niveau de la station)
    AparamOptim        : liste des paramètres à optimiser
                         (les autres restent fixés et lus dans le fichier des stations de calage)
    """

    # initialisations

    # paramètres à optimiser et paramètres fixés
    xparam = [int(x) for x in AdonneesStation['Xparam']]

    # niveau initial de remplissage des réservoirs S/X1 et R/X3
    LniveauIni = np.array([0.6,0.7])

    dt_station = [analyseChroniques.decoder_date(tps) for tps in AdonneesStation['duree_eval']]

    # chronique de pluie (en mm)
    LPluie = np.array(AchroniquesPluie['valeur'])
    # chronique d'évapotranspiration (en mm)
    Levapo = np.array(AchroniquesEvapo['valeur'])

    # chronique des débits mesurés en ayant retiré les valeurs négatives pour la simulation
    LQ_calage = np.array([np.max([val,0]) for val in AchroniquesQmesure['valeur']])
    # intervalle de temps d'évaluation du critère d'efficacité (indice dans le tableau des chroniques de pluie)
    tab_dtPluie = AchroniquesPluie['temps']
    dateP0 = tab_dtPluie[0]
    dateP1 = tab_dtPluie[1]
    dtPluie = dateP1-dateP0
    Ldt_eval = np.array([tab_dtPluie.index(dt_station[0]),tab_dtPluie.index(dt_station[1])])

    # test de validité de l'intervalle d'évaluation
    if np.min(LQ_calage[Ldt_eval[0]:Ldt_eval[1]])<0:
        raise ValueError('Tentative de calage sur un intervalle manquant de donnees : '+str(dt_station))
    # optimisation
    LniveauIni = np.array([0.6,0.7])

##    fichierLog.write('dt_station = '+str(dt_station)+'\r\n')
##    fichierLog.write('Ldt_eval = ' + str(Ldt_eval)+'\r\n')
##    fichierLog.write('mean(LQ_calage) = ' + str(np.mean(LQ_calage[Ldt_eval[0]:Ldt_eval[1]]))+'\r\n')
##    fichierLog.write('min(LQ_calage) = ' + str(np.min(LQ_calage[Ldt_eval[0]:Ldt_eval[1]]))+'\r\n')
##    fichierLog.write('max(LQ_calage) = ' + str(np.max(LQ_calage[Ldt_eval[0]:Ldt_eval[1]]))+'\r\n')

    # 2 cas possibles : pas de temps mensuel -> GR2M, pas de temps inférieur à 30j -> GR4J
    if dtPluie.days == 31:
        # paramètres de GR2M
        Lxini=np.array(xparam) # X1, X2 et X3 à optimiser

        fichierLog.write('Simulation avec GR2M \r\n')
        fichierLog.write('param à optimiser : ' + str(AparamOptim)+'\r\n')
        fichierLog.write('valeurs initiales : ' + str(Lxini)+'\r\n')

        LXfmin = opt.fmin(fmin_gr2m_station,
                          Lxini,
                          (AparamOptim,Lxini,LniveauIni,Ldt_eval,LPluie,Levapo,LQ_calage,dtPluie,
                          AcritereEfficacite,fichierLog),
                          maxiter=1000)
        # résultat
        resultat = LXfmin[:]
    else:
        raise NotImplementedError("GR4J n'est pas implémenté.")

    # paramètres fixés à leur valeur initiale
    if not 'X1' in AparamOptim:
        resultat[0]=Lxini[0]
    if not 'X2' in AparamOptim:
        resultat[1]=Lxini[1]
    if not 'X3' in AparamOptim:
        resultat[2]=Lxini[2]

    fichierLog.write('optimisation terminée\r\n')
    fichierLog.write('X optimal = ' + str(resultat)+'\r\n')

    # paramètres GR2M optimisés et fixés
    return resultat

#-------------------------------------------------------------------------------

def fmin_gr2m_station(x,AparamOptim,Axini,AniveauIni,Adt_eval,APluie,Aevapo,AQ_calage,AdtPluie,AcritereEfficacite,fichierLog):
    """
    fonction de calcul de l'erreur du débit par GR2M par rapport aux données Qj
    lors d'une simulation
    Les paramètre de GR2M sont répartis de la manière suivante:
    x # X1, X2, X3 à optimiser, X4 ignoré
    AparamOptim = liste des paramètres à réellement optimiser,
                  celui non indiqué est laissé fixe à la valeur donnée par Axini

    paramètres :
    x           = variable à optimiser : paramètres de GR2M spécifiés par AparmOptim
    AparamOptim = liste des paramètres à optimiser
    Axini       = valeur initiale des paramètres avant optimisation
    AniveauIni  = vecteur contenant les niveaux de remplissage initiaux des réservoirs (en %) de GR2M
    Adt_eval    = intervalle de temps de simulation pour le calcul de l'erreur
    APluie      = chronique de pluie
    Aevapo      = chronique d'évapotranspiration
    AQ_calage   = chronique de débit mesuré pour calcul d'erreur
    AdtPluie    = pas de temps des chroniques (en jours)
    AcritereEfficacite = critère d'éfficacité choisi pour le calcul de l'erreur
    """

    # initialisation pour simulation GR2M
    lX=np.array(x)
    # paramètres fixés : les paramètres à optimiser sont dans AparamOptim, les autres sont fixées à leur valeur initiale
    if not 'X1' in AparamOptim:
        lX[0]=Axini[0]
    if not 'X2' in AparamOptim:
        lX[1]=Axini[1]
    if not 'X3' in AparamOptim:
        lX[2]=Axini[2]
    # paramètre X4 ignoré : valeur arbitraire
    lX[3]=Axini[3]

    # pas de temps en heures
    Lpas=24*AdtPluie.days

##    fichierLog.write('dans fmin\r\n')

##    fichierLog.write('lX = ' + str(lX)+'\r\n')
##    fichierLog.write('APluie = '+str(APluie[:10])+'\r\n')
##    fichierLog.write('Aevapo = '+str(Aevapo[:10])+'\r\n')
##    fichierLog.write('AniveauIni = '+str(AniveauIni)+'\r\n')
##    fichierLog.write('Lpas = '+str(Lpas)+'\r\n')

    # appel de GR2M
    resultatGR2M = gr2m.gr2m(lX,APluie,Aevapo,AniveauIni,Lpas)

    # calcul de l'erreur selon le critère d'efficacité AcritereEfficacite
    # pas de transformation
    LQ_GR2M = resultatGR2M['Q'][Adt_eval[0]:Adt_eval[1]]
    LQ_ref = AQ_calage[Adt_eval[0]:Adt_eval[1]]

##    fichierLog.write('LQ_ref = ' + str(LQ_ref[:10])+'\r\n')
##    fichierLog.write('mean(LQ_ref) = ' + str(np.mean(LQ_ref))+'\r\n')
##    fichierLog.write('min(LQ_ref) = ' + str(np.min(LQ_ref))+'\r\n')
##    fichierLog.write('max(LQ_ref) = ' + str(np.max(LQ_ref))+'\r\n')


    Nash = calculerNash(LQ_GR2M,LQ_ref,AcritereEfficacite)

    erreur = 1-Nash

    # vérification des bornes de x : erreur augmentée pour rester dans les limites imposées
    if (x[0]<1)|(x[1]<0)|(x[1]>50)|(x[2]<1):
        erreur=erreur*10

    # paramètres à ne pas optimiser : erreur augmentée
    if (not 'X1' in AparamOptim)&(lX[0]!=Axini[0]): erreur=erreur*10
    if (not 'X2' in AparamOptim)&(lX[1]!=Axini[1]): erreur=erreur*10
    if (not 'X3' in AparamOptim)&(lX[2]!=Axini[2]): erreur=erreur*10
    if lX[3]!=Axini[3]: erreur=erreur*10

##    fichierLog.write('Nash = '+str(Nash)+'\r\n')
##    fichierLog.write('erreur = '+str(erreur)+'\r\n')

    # résultat
    return erreur

#-------------------------------------------------------------------------------

def calculerNash(AQ_GR,AQ_ref,AcritereEfficacite):
    """
    calcul du critère de Nash selon le paramètre AcritereEfficacite.
    AQ_GR est la chronique de débit simulé
    AQ_ref est celle de débit mesuré
    Renvoie le critère de Nash calculé / valeurs non nulles de débit mesuré
    """
    # exclusion des valeurs de AQ_ref (les simulations ne génèrent pas de valeurs nulles)
    id_ok=np.where(AQ_ref>0)
    # Q
    if len(id_ok[0]>0):
        lQ_GR = np.array(AQ_GR[id_ok[0]])
        lQ_ref = np.array(AQ_ref[id_ok[0]])
    else:
        lQ_GR = np.array(AQ_GR)
        lQ_ref = np.array(AQ_ref)
    QmesMoyen=np.mean(lQ_ref)
    # sqrt(Q)
    if AcritereEfficacite == typeCritereNash.listeCritereNash.index(typeCritereNash.__NASHVQ__):
        lQ_GR = np.sqrt(lQ_GR)
        lQ_ref = np.sqrt(lQ_ref)
    # log(Q+M/40)
    elif AcritereEfficacite == typeCritereNash.listeCritereNash.index(typeCritereNash.__NASHLOGQ__):
        lQ_GR = np.log(lQ_GR+QmesMoyen/40.)
        lQ_ref = np.log(lQ_ref+QmesMoyen/40.)

    # calcul du critère de Nash
    QrefMoyen=np.mean(lQ_ref)
    Nash=1-(LA.norm(lQ_GR-lQ_ref,2)**2)/(LA.norm(lQ_ref-QrefMoyen,2)**2)
    # résultat
    return Nash

#-------------------------------------------------------------------------------
#-------------------------------------------------------------------------------
